import torch
import torchvision.datasets
from torch import nn
from torch.nn import MaxPool2d
from torch.nn.functional import max_pool2d
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter


class MyMaxPool(nn.Module):
    """最大池化 保持输入的特征 同时降低输入量"""

    def __init__(self):
        super().__init__()
        self.maxpool = MaxPool2d(kernel_size=3, ceil_mode=False)

    def forward(self, input):
        output = self.maxpool(input)
        return output


# # 最大池化操作
# input = torch.tensor([[1, 2, 0, 3, 1],
#                       [0, 1, 2, 3, 1],
#                       [1, 2, 1, 0, 0],
#                       [5, 2, 3, 1, 1],
#                       [2, 1, 0, 1, 1]], dtype=torch.float32)
# input = torch.reshape(input, (1, 1, 5, 5))
# output = max_pool2d(input, 3, ceil_mode=True)
# print(output)
#
# # 最大池化模型
# maxpool1 = MyMaxPool()
# output1 = maxpool1.forward(input)
# print(output1)

# 使用图片验证
dataset = torchvision.datasets.CIFAR10("../datasets/CIFAR10", train=False, transform=torchvision.transforms.ToTensor(),
                                       download=True)
dataloader = DataLoader(dataset, batch_size=64)
writer = SummaryWriter("../logs/cvlogs")
step = 0
maxpool2 = MyMaxPool()

for data in dataloader:
    image, target = data
    writer.add_images("input", image, step)
    output2 = maxpool2(image)
    writer.add_images("output", output2, step)
    step += 1

writer.close()